import base64
from Crypto.Cipher import AES
import sys
import re
import hashlib


class AES128_UDS:
    def __init__(self, Learmask=None) -> None:
        if Learmask is None:
            self.key = 0x01020304050607080102030405060708.to_bytes(length=16)
        else:
            self.key = Learmask

        self.mode = AES.MODE_ECB

    def pad(self, Seed=None):
        '''
        尝试的pkcs5补齐方法
        '''

        if Seed is None:
            self.data = 0x6816b4d5.to_bytes(length=4)
        else:
            self.data = Seed
        count = len(self.data)

        add_count = 16 - count % 16
        pad_item = add_count
        add_data = pad_item.to_bytes(1) * add_count

        # print(f'[file:{__file__}]=>[line:{
        #     (sys._getframe().f_lineno)}]==>add_data{add_data}')
        dataseed = self.data + add_data
        # print(f'[file:{__file__}]=>[line:{
        #     (sys._getframe().f_lineno)}]==>data{dataseed}')

        return dataseed

    def AES_en(self, data):
        # 创建加密对象
        AES_obj = AES.new(self.key, self.mode)

        # 补齐加密数据
        pad_data = self.pad(data)
        # 完成加密
        AES_en_byte = AES_obj.encrypt(pad_data)

        # 转为hex字符串
        AES_en_str = AES_en_byte.hex()

        return AES_en_byte

    def AES_de(self, data):
        '''
            待实现
        '''
        pass


def seed_to_key(positive_resp_data, Learmask="5f0f50a1") -> list[int]:
    # seedValue = [int(item, 16) for item in seed.split(' ')[1:]]

    # seed = [int(item, 16) for item in positive_resp_data.split(' ')[1:]]

    # 将字符串的Learmask转换为32位二进制数
    Learmask_hex = int(Learmask, 16)
    seed = positive_resp_data[1:]
    pass
    # print(f'[file:{__file__}]=>[line:{
    #     (sys._getframe().f_lineno)}]==>seed:{seed}')

    seedValue = (seed[0] << 24) | (
        seed[1] << 16) | (seed[2] << 8) | (seed[3])
    # print("seedValue:", hex(seedValue))

    if seedValue != 0:
        for i in range(35):
            if (seedValue & 0x80000000):
                seedValue = seedValue << 1
                seedValue = seedValue ^ Learmask_hex
            else:
                seedValue = seedValue << 1

        keyValue = seedValue
    else:
        keyValue = 0

    keyValue = keyValue & 0xffffffff
    # print("keyValue:", hex(keyValue))

    key = ()
    for i in range(4):
        key = [keyValue & 0x000000ff,  *key]
        keyValue = keyValue >> 8
    # print("key_hex:", [hex(item) for item in key])

    # 反向排序尝试
    # key = key[::-1]

    return key


def seed_to_key_Qbaic(positive_resp_data=None, Learmask=None):
    '''
        QBAIC_C230089-2021 UDS安全访问认证技术要求
        Learmask: 16字节密匙因子,16进制表示 0x********
        seed: ECU接收到解锁请求后返回的随机种子,4字节
    '''

    if Learmask is None:
        Learmask = 0x7720E9C394E2076F7023D2EF41EAB880.to_bytes(length=16)
    else:

        # 先将Learmask转换为byte字节串
        Learmask = int(Learmask, 16)
        Learmask = Learmask.to_bytes(length=16)
    # print(Learmask)

    # 截取后4个数据
    seed = positive_resp_data[1:]

    # [ ] 1 使用4字节seed通过SHA256算法产生32字节hash值 seedHash256
    # 随机数，直接用16进制表示
    if seed is None:
        seed = 0x21c43252

    pass
    # print(f'[file:{__file__}]=>[line:{
    #     (sys._getframe().f_lineno)}]==>seed:{seed}')

    seedValue = (seed[0] << 24) | (
        seed[1] << 16) | (seed[2] << 8) | (seed[3])

    # seed转换为byte类型
    seed_byte = seedValue.to_bytes(length=4)
    # print(seed_byte)
    #
    seed_hash01 = hashlib.sha256(seed_byte)
    seed_hash01_byte = seed_hash01.digest()
    seed_hash01_str = seed_hash01.hexdigest()
    # print(seed_hash01_str)
    # print(seed_hash01_byte)

    # [ ] 2 拼接seedHash256与Learmask，对拼接的值再进行SHA256算法产生32字节Hash值，取前16位作为key_AES128
    seed_hash01_Learmask = seed_hash01_byte+Learmask

    seed_hash02 = hashlib.sha256(seed_hash01_Learmask)
    seed_hash02_str = seed_hash02.hexdigest()
    seed_hash02_byte = seed_hash02.digest()

    # print(seed_hash02_str)
    # print(seed_hash02_byte)

    # [ ] 3 获取seed_hash02_byte的前16位作为密匙k
    seed_key = seed_hash02_byte[:16]
    # print(seed_key)
    (seed_hash02_str[:32])

    # [ ] 4 使用key_AES128作为秘钥，通过AES128算法(AES128-ECB PKCS5填充)，对seed进行加密，得到16字节的key_unlock
    # aes = AESECB(seed_key)
    aes = AES128_UDS(seed_key)
    # print(aes.encrypt(seed_byte))
    key_unlock = aes.AES_en(seed_byte)
    # ic_debugger(key_unlock)
    # ic_debugger(key_unlock.hex())
    # [ ] 6 将key_unlock返回，后续发回给ECU
    return key_unlock


seed_to_key_algorithm_list = {
    "algo1": seed_to_key,
    "algo_Qbaic": seed_to_key_Qbaic
}

if 0:
    # 参考算法实现素材
    def AES_en(key, data):
        # 将长度不足16字节的字符串补齐
        if len(data) < 16:
            data = pad(data)
        # 创建加密对象
        # AES_obj = AES.new(key.encode("utf-8"), AES.MODE_CBC, iv.encode("utf-8"))
        AES_obj = AES.new(key.encode("utf-8"), AES.MODE_ECB)
        # 完成加密
        AES_en_str = AES_obj.encrypt(data)
        # 用base64编码一下
        AES_en_str = base64.b64encode(AES_en_str)
        # 最后将密文转化成字符串
        AES_en_str = AES_en_str.decode("utf-8")
        return AES_en_str

    def AES_de(key, data):
        # 解密过程逆着加密过程写
        # 将密文字符串重新编码成二进制形式
        data = data.encode("utf-8")
        # 将base64的编码解开
        data = base64.b64decode(data)
        # 创建解密对象
        # AES_de_obj = AES.new(key.encode("utf-8"), AES.MODE_CBC, iv.encode("utf-8"))
        AES_de_obj = AES.new(key.encode("utf-8"), AES.MODE_ECB)
        # 完成解密
        AES_de_str = AES_de_obj.decrypt(data)
        # 去掉补上的空格
        AES_de_str = AES_de_str.strip()
        # 对明文解码
        AES_de_str = AES_de_str.decode("utf-8")
        return AES_de_str

    key = 'miyaoxuyao16ziji'
    data = b'hello world'
    # 将原始的明文用空格填充到16字节
    data = AES_en(key, data)
    print(data)
    data = AES_de(key, data)
    print(data)

    class AESECB:
        '''
            AESECB算法示例类
            [ ] 该类填充可能也不是pkcs5，需要进一步确认
        '''

        from Crypto.Cipher import AES

        def __init__(self, key):
            self.key = key
            self.mode = AES.MODE_ECB
            self.bs = 16  # block size
            self.PADDING = lambda s: s + \
                (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)

        def encrypt(self, text):
            generator = AES.new(self.key, self.mode)  # ECB模式无需向量iv
            crypt = generator.encrypt(self.PADDING(text))
            crypted_str = base64.b64encode(crypt)
            result = crypted_str.decode()
            return result

        def decrypt(self, text):
            generator = AES.new(self.key, self.mode)  # ECB模式无需向量iv
            text += (len(text) % 4) * '='
            decrpyt_bytes = base64.b64decode(text)
            meg = generator.decrypt(decrpyt_bytes)
            # 去除解码后的非法字符
            try:
                result = re.compile(
                    '[\\x00-\\x08\\x0b-\\x0c\\x0e-\\x1f\n\r\t]').sub('', meg.decode())
            except Exception:
                result = '解码失败，请重试!'
            return result

    class EncryptDate2:
        '''
            备选第二个AEB类
            该类用的是pkcs7填充，不能直接使用
        '''

        import base64
        from Crypto.Cipher import AES
        from Crypto.Util.Padding import pad

        def __init__(self, key):
            # 初始化密钥
            self.key = key
            # 初始化数据块大小
            self.length = AES.block_size
            # 初始化AES,ECB模式的实例
            self.aes = AES.new(self.key.encode("utf-8"), AES.MODE_ECB)
            # 截断函数，去除填充的字符
            self.unpad = lambda date: date[0:-ord(date[-1])]

        def fill_method(self, aes_str):
            '''pkcs7补全'''
            pad_pkcs7 = pad(aes_str.encode('utf-8'),
                            AES.block_size, style='pkcs7')

            return pad_pkcs7

        def encrypt(self, encrData):
            # 加密函数,使用pkcs7补全
            res = self.aes.encrypt(self.fill_method(encrData))
            # 转换为base64
            msg = str(base64.b64encode(res), encoding="utf-8")

            return msg

        def decrypt(self, decrData):
            # base64解码
            res = base64.decodebytes(decrData.encode("utf-8"))
            # 解密函数
            msg = self.aes.decrypt(res).decode("utf-8")

            return self.unpad(msg)

    if 0:
        # key的长度需要补长(16倍数),补全方式根据情况而定,未补齐会报错
        # key字符长度决定加密结果,长度16：加密结果AES(128),长度32：结果就是AES(256)
        eg = EncryptDate2("asdfghjklqwertyu")
        # 加密字符串长同样需要16倍数：需注意,不过代码中pad()方法里，帮助实现了补全（补全方式就是pkcs7）
        en = eg.encrypt("测试数据")
        de = eg.decrypt(en)
        print(f"加密结果：{en}")
        print(f"解密结果：{de}")

    # pad1方法
    def pad1(data):
        pad_data = data
        for i in range(0, 16-len(data)):
            pad_data = pad_data + b' '
        return pad_data


def main():
    print("hello world")
    seed = [0x01, 0x94, 0xa4, 0x92, 0x52]
    Learmask = "0x7720E9C394E2076F7023D2EF41EAB880"
    key = seed_to_key_Qbaic(seed, Learmask=Learmask)
    print(key.hex())


if __name__ == '__main__':
    main()
